"""
安全认证服务

提供JWT令牌认证、基于角色的访问控制和API密钥管理，包括：
- JWT令牌生成和验证
- 基于角色的访问控制(RBAC)
- API密钥管理
- 会话管理
- 安全审计日志
"""
import jwt
import hashlib
import secrets
import time
from typing import Dict, List, Any, Optional, Set
from datetime import datetime, timedelta
from dataclasses import dataclass, field
from enum import Enum
import json
import logging

# 配置日志
logging.basicConfig(level=logging.INFO)
security_logger = logging.getLogger('security')


class UserRole(Enum):
    """用户角色枚举"""
    ADMIN = "admin"
    RESEARCHER = "researcher"
    STUDENT = "student"
    GUEST = "guest"


class Permission(Enum):
    """权限枚举"""
    # 研究相关权限
    CREATE_RESEARCH = "create_research"
    READ_RESEARCH = "read_research"
    UPDATE_RESEARCH = "update_research"
    DELETE_RESEARCH = "delete_research"
    
    # 数据相关权限
    ACCESS_DATA = "access_data"
    EXPORT_DATA = "export_data"
    IMPORT_DATA = "import_data"
    
    # 系统管理权限
    MANAGE_USERS = "manage_users"
    MANAGE_SYSTEM = "manage_system"
    VIEW_LOGS = "view_logs"
    
    # 协作权限
    CREATE_COLLABORATION = "create_collaboration"
    JOIN_COLLABORATION = "join_collaboration"
    MANAGE_COLLABORATION = "manage_collaboration"


@dataclass
class User:
    """用户信息"""
    user_id: str
    username: str
    email: str
    password_hash: str
    role: UserRole
    permissions: Set[Permission] = field(default_factory=set)
    created_at: datetime = field(default_factory=datetime.now)
    last_login: Optional[datetime] = None
    is_active: bool = True
    failed_login_attempts: int = 0
    locked_until: Optional[datetime] = None
    metadata: Dict[str, Any] = field(default_factory=dict)


@dataclass
class APIKey:
    """API密钥"""
    key_id: str
    key_hash: str
    user_id: str
    name: str
    permissions: Set[Permission] = field(default_factory=set)
    created_at: datetime = field(default_factory=datetime.now)
    expires_at: Optional[datetime] = None
    last_used: Optional[datetime] = None
    is_active: bool = True
    usage_count: int = 0


@dataclass
class Session:
    """用户会话"""
    session_id: str
    user_id: str
    created_at: datetime = field(default_factory=datetime.now)
    expires_at: datetime = field(default_factory=lambda: datetime.now() + timedelta(hours=24))
    last_activity: datetime = field(default_factory=datetime.now)
    ip_address: Optional[str] = None
    user_agent: Optional[str] = None
    is_active: bool = True


@dataclass
class SecurityEvent:
    """安全事件"""
    event_id: str
    event_type: str
    user_id: Optional[str]
    ip_address: Optional[str]
    timestamp: datetime = field(default_factory=datetime.now)
    details: Dict[str, Any] = field(default_factory=dict)
    severity: str = "info"  # info, warning, error, critical


class SecurityService:
    """安全认证服务"""
    
    def __init__(self, secret_key: str, token_expiry_hours: int = 24):
        self.secret_key = secret_key
        self.token_expiry_hours = token_expiry_hours
        
        # 存储
        self.users: Dict[str, User] = {}
        self.api_keys: Dict[str, APIKey] = {}
        self.sessions: Dict[str, Session] = {}
        self.security_events: List[SecurityEvent] = []
        
        # 角色权限映射
        self.role_permissions = {
            UserRole.ADMIN: {
                Permission.CREATE_RESEARCH, Permission.READ_RESEARCH, 
                Permission.UPDATE_RESEARCH, Permission.DELETE_RESEARCH,
                Permission.ACCESS_DATA, Permission.EXPORT_DATA, Permission.IMPORT_DATA,
                Permission.MANAGE_USERS, Permission.MANAGE_SYSTEM, Permission.VIEW_LOGS,
                Permission.CREATE_COLLABORATION, Permission.JOIN_COLLABORATION, 
                Permission.MANAGE_COLLABORATION
            },
            UserRole.RESEARCHER: {
                Permission.CREATE_RESEARCH, Permission.READ_RESEARCH, 
                Permission.UPDATE_RESEARCH, Permission.DELETE_RESEARCH,
                Permission.ACCESS_DATA, Permission.EXPORT_DATA,
                Permission.CREATE_COLLABORATION, Permission.JOIN_COLLABORATION
            },
            UserRole.STUDENT: {
                Permission.READ_RESEARCH, Permission.UPDATE_RESEARCH,
                Permission.ACCESS_DATA, Permission.JOIN_COLLABORATION
            },
            UserRole.GUEST: {
                Permission.READ_RESEARCH
            }
        }
        
        # 安全配置
        self.security_config = {
            'max_failed_login_attempts': 5,
            'account_lockout_duration_minutes': 30,
            'password_min_length': 8,
            'require_password_complexity': True,
            'session_timeout_hours': 24,
            'api_key_default_expiry_days': 90,
            'enable_audit_logging': True
        }
        
        # 初始化管理员用户
        self._create_default_admin()
    
    def _create_default_admin(self):
        """创建默认管理员用户"""
        admin_id = "admin_001"
        if admin_id not in self.users:
            admin_user = User(
                user_id=admin_id,
                username="admin",
                email="admin@system.local",
                password_hash=self._hash_password("admin123"),
                role=UserRole.ADMIN,
                permissions=self.role_permissions[UserRole.ADMIN].copy()
            )
            self.users[admin_id] = admin_user
            
            self._log_security_event(
                "user_created",
                user_id=admin_id,
                details={"username": "admin", "role": "admin"},
                severity="info"
            )
    
    def register_user(self, username: str, email: str, password: str, 
                     role: UserRole = UserRole.STUDENT) -> Dict[str, Any]:
        """注册新用户"""
        try:
            # 验证输入
            if not self._validate_username(username):
                return {"success": False, "error": "用户名格式无效"}
            
            if not self._validate_email(email):
                return {"success": False, "error": "邮箱格式无效"}
            
            if not self._validate_password(password):
                return {"success": False, "error": "密码不符合安全要求"}
            
            # 检查用户名和邮箱是否已存在
            if self._username_exists(username):
                return {"success": False, "error": "用户名已存在"}
            
            if self._email_exists(email):
                return {"success": False, "error": "邮箱已被使用"}
            
            # 创建用户
            user_id = f"user_{int(time.time() * 1000)}"
            user = User(
                user_id=user_id,
                username=username,
                email=email,
                password_hash=self._hash_password(password),
                role=role,
                permissions=self.role_permissions[role].copy()
            )
            
            self.users[user_id] = user
            
            # 记录安全事件
            self._log_security_event(
                "user_registered",
                user_id=user_id,
                details={"username": username, "email": email, "role": role.value},
                severity="info"
            )
            
            return {
                "success": True,
                "user_id": user_id,
                "message": "用户注册成功"
            }
            
        except Exception as e:
            self._log_security_event(
                "user_registration_failed",
                details={"username": username, "error": str(e)},
                severity="error"
            )
            return {"success": False, "error": "注册失败"}
    
    def authenticate_user(self, username: str, password: str, 
                         ip_address: Optional[str] = None) -> Dict[str, Any]:
        """用户认证"""
        try:
            # 查找用户
            user = self._find_user_by_username(username)
            if not user:
                self._log_security_event(
                    "login_failed",
                    details={"username": username, "reason": "user_not_found"},
                    ip_address=ip_address,
                    severity="warning"
                )
                return {"success": False, "error": "用户名或密码错误"}
            
            # 检查账户是否被锁定
            if self._is_account_locked(user):
                self._log_security_event(
                    "login_failed",
                    user_id=user.user_id,
                    details={"username": username, "reason": "account_locked"},
                    ip_address=ip_address,
                    severity="warning"
                )
                return {"success": False, "error": "账户已被锁定，请稍后再试"}
            
            # 验证密码
            if not self._verify_password(password, user.password_hash):
                user.failed_login_attempts += 1
                
                # 检查是否需要锁定账户
                if user.failed_login_attempts >= self.security_config['max_failed_login_attempts']:
                    lockout_duration = timedelta(minutes=self.security_config['account_lockout_duration_minutes'])
                    user.locked_until = datetime.now() + lockout_duration
                    
                    self._log_security_event(
                        "account_locked",
                        user_id=user.user_id,
                        details={"username": username, "failed_attempts": user.failed_login_attempts},
                        ip_address=ip_address,
                        severity="error"
                    )
                
                self._log_security_event(
                    "login_failed",
                    user_id=user.user_id,
                    details={"username": username, "reason": "invalid_password"},
                    ip_address=ip_address,
                    severity="warning"
                )
                return {"success": False, "error": "用户名或密码错误"}
            
            # 认证成功
            user.failed_login_attempts = 0
            user.locked_until = None
            user.last_login = datetime.now()
            
            # 生成JWT令牌
            token = self._generate_jwt_token(user)
            
            # 创建会话
            session = self._create_session(user.user_id, ip_address)
            
            self._log_security_event(
                "login_successful",
                user_id=user.user_id,
                details={"username": username},
                ip_address=ip_address,
                severity="info"
            )
            
            return {
                "success": True,
                "token": token,
                "session_id": session.session_id,
                "user": {
                    "user_id": user.user_id,
                    "username": user.username,
                    "email": user.email,
                    "role": user.role.value,
                    "permissions": [p.value for p in user.permissions]
                }
            }
            
        except Exception as e:
            self._log_security_event(
                "authentication_error",
                details={"username": username, "error": str(e)},
                ip_address=ip_address,
                severity="error"
            )
            return {"success": False, "error": "认证过程出现错误"}
    
    def verify_token(self, token: str) -> Dict[str, Any]:
        """验证JWT令牌"""
        try:
            payload = jwt.decode(token, self.secret_key, algorithms=['HS256'])
            
            user_id = payload.get('user_id')
            if not user_id or user_id not in self.users:
                return {"valid": False, "error": "用户不存在"}
            
            user = self.users[user_id]
            if not user.is_active:
                return {"valid": False, "error": "用户已被禁用"}
            
            return {
                "valid": True,
                "user": {
                    "user_id": user.user_id,
                    "username": user.username,
                    "role": user.role.value,
                    "permissions": [p.value for p in user.permissions]
                }
            }
            
        except jwt.ExpiredSignatureError:
            return {"valid": False, "error": "令牌已过期"}
        except jwt.InvalidTokenError:
            return {"valid": False, "error": "无效令牌"}
        except Exception as e:
            return {"valid": False, "error": "令牌验证失败"}
    
    def check_permission(self, user_id: str, permission: Permission) -> bool:
        """检查用户权限"""
        if user_id not in self.users:
            return False
        
        user = self.users[user_id]
        return permission in user.permissions
    
    def create_api_key(self, user_id: str, name: str, 
                      permissions: Optional[Set[Permission]] = None,
                      expires_days: Optional[int] = None) -> Dict[str, Any]:
        """创建API密钥"""
        try:
            if user_id not in self.users:
                return {"success": False, "error": "用户不存在"}
            
            user = self.users[user_id]
            
            # 生成API密钥
            api_key = secrets.token_urlsafe(32)
            key_id = f"key_{int(time.time() * 1000)}"
            
            # 设置权限（默认继承用户权限）
            if permissions is None:
                permissions = user.permissions.copy()
            else:
                # 确保不超过用户权限
                permissions = permissions.intersection(user.permissions)
            
            # 设置过期时间
            expires_at = None
            if expires_days:
                expires_at = datetime.now() + timedelta(days=expires_days)
            elif self.security_config['api_key_default_expiry_days']:
                expires_at = datetime.now() + timedelta(
                    days=self.security_config['api_key_default_expiry_days']
                )
            
            # 创建API密钥记录
            api_key_record = APIKey(
                key_id=key_id,
                key_hash=self._hash_api_key(api_key),
                user_id=user_id,
                name=name,
                permissions=permissions,
                expires_at=expires_at
            )
            
            self.api_keys[key_id] = api_key_record
            
            self._log_security_event(
                "api_key_created",
                user_id=user_id,
                details={"key_id": key_id, "name": name},
                severity="info"
            )
            
            return {
                "success": True,
                "api_key": api_key,
                "key_id": key_id,
                "expires_at": expires_at.isoformat() if expires_at else None
            }
            
        except Exception as e:
            self._log_security_event(
                "api_key_creation_failed",
                user_id=user_id,
                details={"name": name, "error": str(e)},
                severity="error"
            )
            return {"success": False, "error": "API密钥创建失败"}
    
    def verify_api_key(self, api_key: str) -> Dict[str, Any]:
        """验证API密钥"""
        try:
            key_hash = self._hash_api_key(api_key)
            
            # 查找匹配的API密钥
            for key_record in self.api_keys.values():
                if key_record.key_hash == key_hash and key_record.is_active:
                    # 检查是否过期
                    if key_record.expires_at and datetime.now() > key_record.expires_at:
                        return {"valid": False, "error": "API密钥已过期"}
                    
                    # 更新使用记录
                    key_record.last_used = datetime.now()
                    key_record.usage_count += 1
                    
                    # 获取用户信息
                    user = self.users.get(key_record.user_id)
                    if not user or not user.is_active:
                        return {"valid": False, "error": "关联用户无效"}
                    
                    return {
                        "valid": True,
                        "key_id": key_record.key_id,
                        "user_id": key_record.user_id,
                        "permissions": [p.value for p in key_record.permissions]
                    }
            
            return {"valid": False, "error": "无效的API密钥"}
            
        except Exception as e:
            return {"valid": False, "error": "API密钥验证失败"}
    
    def revoke_api_key(self, user_id: str, key_id: str) -> Dict[str, Any]:
        """撤销API密钥"""
        try:
            if key_id not in self.api_keys:
                return {"success": False, "error": "API密钥不存在"}
            
            key_record = self.api_keys[key_id]
            
            # 检查权限（只能撤销自己的密钥或管理员可以撤销任何密钥）
            user = self.users.get(user_id)
            if not user:
                return {"success": False, "error": "用户不存在"}
            
            if key_record.user_id != user_id and user.role != UserRole.ADMIN:
                return {"success": False, "error": "权限不足"}
            
            # 撤销密钥
            key_record.is_active = False
            
            self._log_security_event(
                "api_key_revoked",
                user_id=user_id,
                details={"key_id": key_id, "revoked_by": user_id},
                severity="info"
            )
            
            return {"success": True, "message": "API密钥已撤销"}
            
        except Exception as e:
            return {"success": False, "error": "撤销API密钥失败"}
    
    def logout_user(self, session_id: str) -> Dict[str, Any]:
        """用户登出"""
        try:
            if session_id not in self.sessions:
                return {"success": False, "error": "会话不存在"}
            
            session = self.sessions[session_id]
            session.is_active = False
            
            self._log_security_event(
                "user_logout",
                user_id=session.user_id,
                details={"session_id": session_id},
                severity="info"
            )
            
            return {"success": True, "message": "登出成功"}
            
        except Exception as e:
            return {"success": False, "error": "登出失败"}
    
    def _generate_jwt_token(self, user: User) -> str:
        """生成JWT令牌"""
        payload = {
            'user_id': user.user_id,
            'username': user.username,
            'role': user.role.value,
            'iat': datetime.now(),
            'exp': datetime.now() + timedelta(hours=self.token_expiry_hours)
        }
        
        return jwt.encode(payload, self.secret_key, algorithm='HS256')
    
    def _create_session(self, user_id: str, ip_address: Optional[str] = None) -> Session:
        """创建用户会话"""
        session_id = f"session_{secrets.token_urlsafe(16)}"
        session = Session(
            session_id=session_id,
            user_id=user_id,
            ip_address=ip_address
        )
        
        self.sessions[session_id] = session
        return session
    
    def _hash_password(self, password: str) -> str:
        """哈希密码"""
        salt = secrets.token_hex(16)
        password_hash = hashlib.pbkdf2_hmac('sha256', password.encode(), salt.encode(), 100000)
        return f"{salt}:{password_hash.hex()}"
    
    def _verify_password(self, password: str, password_hash: str) -> bool:
        """验证密码"""
        try:
            salt, hash_hex = password_hash.split(':')
            password_hash_bytes = hashlib.pbkdf2_hmac('sha256', password.encode(), salt.encode(), 100000)
            return password_hash_bytes.hex() == hash_hex
        except:
            return False
    
    def _hash_api_key(self, api_key: str) -> str:
        """哈希API密钥"""
        return hashlib.sha256(api_key.encode()).hexdigest()
    
    def _validate_username(self, username: str) -> bool:
        """验证用户名"""
        if not username or len(username) < 3 or len(username) > 50:
            return False
        
        # 只允许字母、数字和下划线
        import re
        return re.match(r'^[a-zA-Z0-9_]+$', username) is not None
    
    def _validate_email(self, email: str) -> bool:
        """验证邮箱"""
        import re
        pattern = r'^[a-zA-Z0-9._%+-]+@[a-zA-Z0-9.-]+\.[a-zA-Z]{2,}$'
        return re.match(pattern, email) is not None
    
    def _validate_password(self, password: str) -> bool:
        """验证密码强度"""
        if len(password) < self.security_config['password_min_length']:
            return False
        
        if self.security_config['require_password_complexity']:
            # 要求包含大小写字母、数字和特殊字符
            import re
            if not re.search(r'[a-z]', password):
                return False
            if not re.search(r'[A-Z]', password):
                return False
            if not re.search(r'\d', password):
                return False
            if not re.search(r'[!@#$%^&*(),.?":{}|<>]', password):
                return False
        
        return True
    
    def _username_exists(self, username: str) -> bool:
        """检查用户名是否存在"""
        return any(user.username == username for user in self.users.values())
    
    def _email_exists(self, email: str) -> bool:
        """检查邮箱是否存在"""
        return any(user.email == email for user in self.users.values())
    
    def _find_user_by_username(self, username: str) -> Optional[User]:
        """根据用户名查找用户"""
        for user in self.users.values():
            if user.username == username:
                return user
        return None
    
    def _is_account_locked(self, user: User) -> bool:
        """检查账户是否被锁定"""
        if user.locked_until is None:
            return False
        return datetime.now() < user.locked_until
    
    def _log_security_event(self, event_type: str, user_id: Optional[str] = None,
                          ip_address: Optional[str] = None, details: Optional[Dict[str, Any]] = None,
                          severity: str = "info"):
        """记录安全事件"""
        if not self.security_config['enable_audit_logging']:
            return
        
        event = SecurityEvent(
            event_id=f"event_{int(time.time() * 1000)}",
            event_type=event_type,
            user_id=user_id,
            ip_address=ip_address,
            details=details or {},
            severity=severity
        )
        
        self.security_events.append(event)
        
        # 记录到日志
        security_logger.info(f"Security Event: {event_type} - User: {user_id} - IP: {ip_address} - Details: {details}")
        
        # 限制事件历史大小
        if len(self.security_events) > 10000:
            self.security_events = self.security_events[-5000:]
    
    def get_security_events(self, user_id: Optional[str] = None, 
                          event_type: Optional[str] = None,
                          limit: int = 100) -> List[Dict[str, Any]]:
        """获取安全事件"""
        events = self.security_events
        
        # 过滤条件
        if user_id:
            events = [e for e in events if e.user_id == user_id]
        
        if event_type:
            events = [e for e in events if e.event_type == event_type]
        
        # 按时间倒序排列
        events = sorted(events, key=lambda x: x.timestamp, reverse=True)
        
        # 限制数量
        events = events[:limit]
        
        # 转换为字典格式
        return [
            {
                'event_id': e.event_id,
                'event_type': e.event_type,
                'user_id': e.user_id,
                'ip_address': e.ip_address,
                'timestamp': e.timestamp.isoformat(),
                'details': e.details,
                'severity': e.severity
            }
            for e in events
        ]
    
    def get_user_info(self, user_id: str) -> Optional[Dict[str, Any]]:
        """获取用户信息"""
        if user_id not in self.users:
            return None
        
        user = self.users[user_id]
        return {
            'user_id': user.user_id,
            'username': user.username,
            'email': user.email,
            'role': user.role.value,
            'permissions': [p.value for p in user.permissions],
            'created_at': user.created_at.isoformat(),
            'last_login': user.last_login.isoformat() if user.last_login else None,
            'is_active': user.is_active,
            'failed_login_attempts': user.failed_login_attempts,
            'locked_until': user.locked_until.isoformat() if user.locked_until else None
        }
    
    def list_api_keys(self, user_id: str) -> List[Dict[str, Any]]:
        """列出用户的API密钥"""
        user_keys = [
            key for key in self.api_keys.values()
            if key.user_id == user_id
        ]
        
        return [
            {
                'key_id': key.key_id,
                'name': key.name,
                'permissions': [p.value for p in key.permissions],
                'created_at': key.created_at.isoformat(),
                'expires_at': key.expires_at.isoformat() if key.expires_at else None,
                'last_used': key.last_used.isoformat() if key.last_used else None,
                'is_active': key.is_active,
                'usage_count': key.usage_count
            }
            for key in user_keys
        ]
    
    def get_security_stats(self) -> Dict[str, Any]:
        """获取安全统计"""
        total_users = len(self.users)
        active_users = len([u for u in self.users.values() if u.is_active])
        locked_users = len([u for u in self.users.values() if self._is_account_locked(u)])
        
        total_api_keys = len(self.api_keys)
        active_api_keys = len([k for k in self.api_keys.values() if k.is_active])
        
        active_sessions = len([s for s in self.sessions.values() if s.is_active])
        
        # 最近24小时的安全事件
        recent_events = [
            e for e in self.security_events
            if (datetime.now() - e.timestamp).total_seconds() < 86400
        ]
        
        event_counts = {}
        for event in recent_events:
            event_counts[event.event_type] = event_counts.get(event.event_type, 0) + 1
        
        return {
            'users': {
                'total': total_users,
                'active': active_users,
                'locked': locked_users
            },
            'api_keys': {
                'total': total_api_keys,
                'active': active_api_keys
            },
            'sessions': {
                'active': active_sessions
            },
            'recent_events': {
                'total': len(recent_events),
                'by_type': event_counts
            },
            'security_config': self.security_config
        }